# -*- coding: utf-8 -*-
"""
将人工核对出的结果合并到一起
人工核对的结果保存在csv文件中，遍历所有csv文件，并合并到一起, 保存在一个文件中: ground_truth.txt
"""

import re
import os
from datetime import datetime
import numpy as np
import oyaml as yaml

station_id_map = {
    19: "B11路 下行 兰陵路光华路",
    20: "B11路 下行 兰陵路劳动路",
    21: "B11路 下行 广化桥",
    36: "B11路 上行 劳动路广化街",
    37: "B11路 上行 兰陵路劳动路",
    38: "B11路 上行 兰陵路光华路",
    39: "B11路 上行 兰陵路中吴大道",
    40: "B11路 上行 兰陵路聚湖路",
    68: "B1路 上行 兰陵路聚湖路",
    70: "B1路 上行 兰陵路光华路",
    71: "B1路 上行 兰陵路劳动路",
    72: "B1路 上行 劳动路广化街",
    103: "B1路 下行 劳动路广化街",
    104: "B1路 下行 兰陵路劳动路",
    106: "B1路 下行 兰陵路中吴大道",
    107: "B1路 下行 兰陵路聚湖路"
}


def read_pkg_file(pkg_file_path):
    """
    以二进制格式读取特征文本文件
    Parameters
    ----------
    pkg_file_path: string
        特征文本文件地址

    Returns
    -------
    dict:
    """
    file = open(pkg_file_path, "rb")
    property_infos = file.read(os.path.getsize(pkg_file_path) - 3072 * 4)
    property_infos = property_infos.split(b'version')[-1]
    property_infos_list = str(property_infos).split("\\r")
    uodr_woker_id = int(property_infos_list[3].split(":")[-1])
    uodr_begin_time = int(property_infos_list[4].split(":")[-1])
    uodr_end_time = int(property_infos_list[5].split(":")[-1])
    uodr_report_time = int(property_infos_list[6].split(":")[-1])
    uodr_report_index = int(property_infos_list[7].split(":")[-1])
    uodr_up_or_down = int(property_infos_list[8].split(":")[-1])
    uodr_station_id = int(property_infos_list[9].split(":")[-1])
    uodr_station_status = int(property_infos_list[10].split(":")[-1])
    uodr_pkgs_size = int(property_infos_list[1].split(":")[-1])
    uodr_pkg = np.frombuffer(file.read(), dtype=np.float32).tolist()
    file.close()
    return {
        "uodr_woker_id": uodr_woker_id,
        "uodr_begin_time": uodr_begin_time,
        "uodr_end_time": uodr_end_time,
        "uodr_report_time": uodr_report_time,
        "uodr_report_index": uodr_report_index,
        "uodr_up_or_down": uodr_up_or_down,
        "uodr_station_id": uodr_station_id,
        "uodr_station_status": uodr_station_status,
        "uodr_pkgs_size": uodr_pkgs_size,
        "uodr_pkg": uodr_pkg
    }


def parse_check_file_element(element, root_folder):
    """
    解析check file中每一行中的gallery和query记录。
    两者解析的方式是相同的，此处使用相同的函数进行处理。
    """
    pattern = re.compile(r".*?(?P<mac>[0-9A-Z]{12})/"
                         r"(?P<camera_id>\d{1})-"
                         r"(?P<timestamp>\d{10})-"
                         r"(?P<person_id>\d{1,3})_"
                         r"(?P<station_id>\d{1,3})")
    match = pattern.search(element)
    assert match is not None, "匹配失败"
    match_dict = match.groupdict()
    mac = match_dict["mac"]
    camera_id = "c" + match_dict["camera_id"]
    timestamp = match_dict["timestamp"]
    person_id = match_dict["person_id"]

    image_folder = os.path.join("upload", mac, camera_id, timestamp, person_id)
    reid_feature_path = os.path.join(
        "reid_pkg",
        datetime.fromtimestamp(int(timestamp)).strftime("%Y-%m-%d"),
        element.rstrip().replace(" ", "").replace("/", "\\"))
    info = read_pkg_file(os.path.join(root_folder, reid_feature_path))
    time = datetime.fromtimestamp(int(
        info["uodr_begin_time"])).strftime("%Y-%m-%d %H:%M:%S")
    location = info["uodr_station_id"]
    feature = info["uodr_pkg"]
    return image_folder, reid_feature_path, time, location, feature


def parse_check_file_line(line, root_folder):
    gallery_element, query_element = line.split(",")
    image_folder, reid_feature_path, time, location, g_feature = parse_check_file_element(
        gallery_element, root_folder)
    gallery = {
        "image": image_folder,
        "reid_feature": reid_feature_path,
        "time": time,
        "location": station_id_map[location]
    }
    image_folder, reid_feature_path, time, location, q_feature = parse_check_file_element(
        query_element, root_folder)
    query = {
        "image": image_folder,
        "reid_feature": reid_feature_path,
        "time": time,
        "location": station_id_map[location]
    }
    distance = 1 - np.dot(np.array(g_feature), np.array(q_feature))
    distance = F"{distance.item():.2f}"
    return dict(gallery=gallery, query=query, distance=distance)


def main():
    csv_file_folder = r"G:\work\跨线OD\log"
    root_folder = r"G:\work\跨线OD\transfer_reid_check_match"

    csv_file_names = [
        name for name in os.listdir(csv_file_folder) if ".csv" in name
    ]

    pairs = dict()
    pair_index = 1
    for name in csv_file_names:
        csv_file_path = os.path.join(csv_file_folder, name)
        with open(csv_file_path) as csv_file:
            for line in csv_file:
                if "gallery" in line or "query" in line:
                    continue
                pair = parse_check_file_line(line, root_folder)
                pairs[F"pair_{pair_index:05}"] = pair
                pair_index += 1
    with open("ground_truth.yaml", "w", encoding="utf-8") as ground_truth_file:
        yaml.dump(
            {"root": root_folder, "pairs": pairs},
            ground_truth_file,
            allow_unicode=True,
            default_flow_style=False)


if __name__ == "__main__":
    main()
